import matplotlib
matplotlib.use('TkAgg')
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from collections import deque
import pickle
import numpy as np
import random
import math

with open("data/data_floors_0427.pkl", "rb") as handle:
    data_floors = pickle.load(handle)

num_floors, floor_width, floor_length, floor_height, exits_on_ground, cell_size, speed_floor_down_radio, time_update, floors, stairs_up, stairs_down, exits, people = data_floors

# 势能场计算（每层）
static_dist = np.full((num_floors, floor_width, floor_length), np.inf)
dq = deque()
for z, (i, j) in exits:
    dq.append((z, i, j))
    static_dist[z, i, j] = 0

while dq:
    cz, ci, cj = dq.popleft()
    if floors[cz, ci, cj] == 4:  # 上楼梯移动到下一层
        nz, ni, nj = cz + 1, ci, cj
        if 0 <= nz < num_floors and floors[nz, ni, nj] in (0, 3, 4, 5):
            move_length = floor_height
            if static_dist[nz, ni, nj] > static_dist[cz, ci, cj] + move_length:
                static_dist[nz, ni, nj] = static_dist[cz, ci, cj] + move_length
                dq.append((nz, ni, nj))
    # 平层移动
    for di, dj in [(-1,0),(1,0),(0,-1),(0,1),(-1,-1),(-1,1),(1,-1),(1,1)]:
        ni, nj = ci+di, cj+dj
        if 0<=ni<floor_width and 0<=nj<floor_length:
            move_length = math.hypot(di, dj)
            if floors[cz, ni, nj] in (0,3,4,5) and static_dist[cz, ni, nj] > static_dist[cz, ci, cj] + move_length:
                static_dist[cz, ni, nj] = static_dist[cz, ci, cj] + move_length
                dq.append((cz, ni, nj))

# 动态模拟部分
fig, axes = plt.subplots(1, num_floors, figsize=(5*num_floors,5))
fig.suptitle("Enhanced Evacuation Simulation")
plt.tight_layout()
time_elapsed = 0

def density_at_target(f, target_pos, people):
    ti, tj = target_pos
    count = 0
    for p in people:
        pf, (pi, pj), _ = p
        if pf == f:
            dist = math.hypot(pi - ti, pj - tj)
            if dist <= 1:
                count += 1
    area = math.pi * (1.5 * cell_size)**2
    return count / area  # 返回密度

def update(_):
    global people, time_elapsed
    time_elapsed += time_update
    
    for ax in axes:
        ax.clear()
    
    # 绘制建筑结构
    for f in range(num_floors):
        ax = axes[f]
        ax.set_title(f'Floor {f+1}')
        ax.set_xticks([]); ax.set_yticks([])
        grid = floors[f]
        ax.imshow(grid==1, cmap='Greys', vmin=0, vmax=1)
        
        ups = np.array(stairs_up[f])
        if len(ups) > 0:
            ax.plot(ups[:,1], ups[:,0], 'b^', markersize=10, label='Upstairs')
        
        downs = np.array(stairs_down[f])
        if len(downs) > 0:
            ax.plot(downs[:,1], downs[:,0], 'gv', markersize=10, label='Downstairs')
        
        if f == 0:
            exs = np.array([pos for _, pos in exits])
            if len(exs) > 0:
                ax.plot(exs[:,1], exs[:,0], 'rs', markersize=10, label='Exit')
    
    new_people = []
    floor_count = [0]*num_floors
    for p in people:
        f, (i,j), move_length = p
        
        if floors[f,i,j] == 3 and f == 0:
            continue  # 已到达出口
        
        # 处理移动中的楼梯
        if move_length > 0 and move_length != math.inf:
            if floors[f,i,j] == 4:  # 正在上楼梯
                new_people.append([f, (i,j), move_length - speed_max * speed_floor_down_radio * time_update])
                floor_count[f+1] +=1
                axes[f+1].plot(j,i,'r.')  # 绘制在上一层
            else:  # 平层移动
                new_people.append([f, (i,j), move_length - speed_max * time_update])
                floor_count[f] +=1
                axes[f].plot(j,i,'r.')  # 绘制在当前层
            continue
        
        # 处理楼梯口
        if floors[f,i,j] == 5:  # 正在下楼梯
            new_people.append([f-1, (i,j), floor_height])  # 移动到下一层
            floor_count[f] +=1
            axes[f].plot(j,i,'r.')  # 绘制在当前层
            continue
        
        # 计算移动方向
        best_score = -np.inf
        best_move = (i,j)
        weight_c1 = 1.0  # 距离权重
        weight_c2 = 1.0  # 密度权重
        weight_c3 = 0.1  # 从众权重
        
        for di, dj in [(-1,0),(1,0),(0,-1),(0,1),(-1,-1),(-1,1),(1,-1),(1,1)]:
            ni, nj = i+di, j+dj
            if 0<=ni<floor_width and 0<=nj<floor_length:
                if floors[f,ni,nj] in (0,3,4,5):
                    current_dist = static_dist[f, ni, nj]
                    target_density = density_at_target(f, (ni, nj), people)
                    support = target_density * (math.pi * (1.5*cell_size)**2)  # 计算周围人数
                    
                    # 计算各因素评分
                    score_c1 = (1/(current_dist + 1e-6)) * weight_c1
                    score_c2 = (1/(target_density + 1e-6)) * weight_c2
                    score_c3 = (support) * weight_c3
                    
                    total_score = score_c1 + score_c2 + score_c3
                    
                    if total_score > best_score:
                        best_score = total_score
                        best_move = (ni, nj)
                    elif total_score == best_score:
                        if random.random() < 0.5:
                            best_move = (ni, nj)
        
        ni, nj = best_move
        move_length = math.hypot(ni - i, nj - j)
        
        new_people.append([f, (ni, nj), move_length])
        floor_count[f] +=1
        axes[f].plot(nj + random.uniform(-0.1,0.1), ni + random.uniform(-0.1,0.1), 'r.')  # 添加随机偏移
        
    people[:] = new_people
    
    # 显示统计信息
    for f in range(num_floors):
        axes[f].text(0.5, -0.15, f'People: {floor_count[f]}', transform=axes[f].transAxes, ha='center', color='blue')
    plt.suptitle(f'Elapsed Time: {time_elapsed:.1f}s', fontsize=16)
    
    if not people:
        print(f"Evacuation completed in {time_elapsed:.1f} seconds.")
        ani.event_source.stop()

ani = animation.FuncAnimation(fig, update, frames=600, interval=500)
plt.show()